package top.zhangjianyong.tools.indicator.strategy.impl;

import org.springframework.stereotype.Component;
import top.zhangjianyong.tools.entity.ETFData;
import top.zhangjianyong.tools.indicator.context.IndicatorContext;
import top.zhangjianyong.tools.indicator.model.TechnicalIndicatorResult;

import java.math.BigDecimal;
import java.math.RoundingMode;
import java.util.List;

/**
 * ATR（平均真实波幅）指标计算器
 * 窗口14
 *
 * @author zhangjianyong
 * @date 2024/01/01
 */
@Component
public class ATRIndicator extends AbstractIndicator {

    private static final String INDICATOR_NAME = "ATR";
    private static final int PERIOD = 14;

    @Override
    public String getName() {
        return INDICATOR_NAME;
    }

    @Override
    public int getMinDataSize() {
        return PERIOD + 1; // 需要至少15个数据点（因为需要计算前一日收盘价）
    }

    @Override
    protected TechnicalIndicatorResult doCalculate(IndicatorContext context) {
        List<ETFData> dataList = context.getDataList();

        // 计算真实波幅序列
        List<BigDecimal> trueRanges = calculateTrueRanges(dataList);
        if (trueRanges.size() < PERIOD) {
            return TechnicalIndicatorResult.failure(getName(), "真实波幅数据不足");
        }

        // 计算ATR（使用Wilder's平滑方法）
        BigDecimal atr = calculateATR(trueRanges);

        TechnicalIndicatorResult result = TechnicalIndicatorResult.success(getName());
        result.addValue("atr", atr);
        result.addValue("atr_14", atr);

        return result;
    }

    /**
     * 计算真实波幅序列
     * 真实波幅 = max(最高价-最低价, |最高价-前收盘价|, |最低价-前收盘价|)
     * 由于ETFData只有value（收盘价），我们使用日增长率来近似计算
     */
    private List<BigDecimal> calculateTrueRanges(List<ETFData> dataList) {
        List<BigDecimal> trueRanges = new java.util.ArrayList<>();

        for (int i = 1; i < dataList.size(); i++) {
            ETFData current = dataList.get(i);
            ETFData previous = dataList.get(i - 1);

            BigDecimal currentPrice = current.getValue();
            BigDecimal previousPrice = previous.getValue();

            // 由于没有最高价和最低价，使用价格变化幅度作为真实波幅的近似
            // 真实波幅 ≈ |当前价格 - 前收盘价|
            BigDecimal trueRange = currentPrice.subtract(previousPrice).abs();

            // 如果日增长率可用，可以更精确地估算
            // 假设日内波动约为日增长率的1.5倍（简化处理）
            if (current.getDailyRate() != null) {
                BigDecimal dailyRateAbs = current.getDailyRate().abs();
                BigDecimal estimatedRange = currentPrice.multiply(dailyRateAbs)
                    .divide(new BigDecimal(100), 4, RoundingMode.HALF_UP)
                    .multiply(new BigDecimal("1.5"));
                trueRange = trueRange.max(estimatedRange);
            }

            trueRanges.add(trueRange);
        }

        return trueRanges;
    }

    /**
     * 使用Wilder's平滑方法计算ATR
     */
    private BigDecimal calculateATR(List<BigDecimal> trueRanges) {
        if (trueRanges.size() < PERIOD) {
            return BigDecimal.ZERO;
        }

        // 计算初始ATR（前PERIOD个真实波幅的平均值）
        BigDecimal initialATR = BigDecimal.ZERO;
        for (int i = 0; i < PERIOD; i++) {
            initialATR = initialATR.add(trueRanges.get(i));
        }
        initialATR = initialATR.divide(new BigDecimal(PERIOD), 4, RoundingMode.HALF_UP);

        // 使用Wilder's平滑公式更新ATR
        BigDecimal currentATR = initialATR;
        BigDecimal smoothFactor = new BigDecimal(PERIOD - 1);
        BigDecimal period = new BigDecimal(PERIOD);

        for (int i = PERIOD; i < trueRanges.size(); i++) {
            BigDecimal trueRange = trueRanges.get(i);
            // ATR = (前一日ATR × (N-1) + 今日真实波幅) / N
            currentATR = currentATR.multiply(smoothFactor)
                .add(trueRange)
                .divide(period, 4, RoundingMode.HALF_UP);
        }

        return currentATR.setScale(4, RoundingMode.HALF_UP);
    }
}

